'''
https://leetcode.cn/problems/find-k-pairs-with-smallest-sums/?envType=study-plan-v2&envId=top-interview-150
'''
import heapq
from typing import List


class Solution:

    # 堆
    # 0,0 一定是最小的， m-1, n-1, 一定是最大的
    # 初始，将i,0 都加入堆中
    # 选择一个i,j后，将 i,j+1加入待选项中
    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
        m, n = len(nums1), len(nums2)
        heap = [(nums1[i] + nums2[0], i, 0) for i in range(min(m, k+1))]
        heapq.heapify(heap)
        res = []
        while len(res) < k:
            _,i,j = heapq.heappop(heap)
            res.append([nums1[i], nums2[j]])
            if j + 1 < n:
                heapq.heappush(heap, (nums1[i] + nums2[j+1], i, j+1))
        return res

    # 二分找到 第k个的target sum, 然后收集
    # 收集时注意， == target sum，不一定都是res
    def kSmallestPairs2(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
        m, n = len(nums1), len(nums2)
        # 找到，sum(对) <= x的 共有多少对
        def f(x):
            # 梯形，不回退，从右下，往(上,[左上])，查找不回退
            i = 0
            j = n-1
            cnt = 0
            while i < m and j >= 0:
                if nums1[i] + nums2[j] <= x:
                    cnt += j + 1
                    i += 1
                else:       # i+1, j更不会满足，所以可以直接跳掉
                    j -= 1
            return cnt
        l, r = nums1[0] + nums2[0], nums1[-1] + nums2[-1]
        k_sum = r
        while l <= r:
            mid = l + r >> 1
            if f(mid) >= k:
                k_sum = mid
                r = mid - 1
            else:
                l = mid + 1

        # 收集 <= k_sum的对
        # print(k_sum)
        # 分开收集，收集 < k_sum，还是梯形收集，不回退
        res = []
        i, j = 0, n-1
        while i < m and j >= 0:
            if nums1[i] + nums2[j] < k_sum:
                res.extend([[nums1[i], nums2[k]]for k in range(j + 1)])
                i += 1
            else:
                j -= 1
        if len(res) == k: return res
        # 分开收集，收集 == k_sum，还是梯形收集，不回退
        # 当 收够k个，直接 return
        i, j = 0, n-1
        while i < m and j >= 0:
            if nums1[i] + nums2[j] < k_sum:
                i += 1
            elif nums1[i] + nums2[j] == k_sum:
                res.append([nums1[i], nums2[j]])
                if len(res) == k: return res
                temp_j = j - 1
                while temp_j >= 0 and nums2[temp_j] == nums2[j]:
                    res.append([nums1[i], nums2[j]])
                    temp_j -= 1
                    if len(res) == k: return res
                i += 1
            else:
                j -= 1

        return res





nums1 = [1,2,4,5,6]
nums2 = [3,5,7,9]
k = 3
print(Solution().kSmallestPairs2(nums1, nums2, k))

